!! Copyright (C) 2009,2010,2011,2012  Marco Restelli
!!
!! This file is part of:
!!   FEMilaro -- Finite Element Method toolkit
!!
!! FEMilaro is free software; you can redistribute it and/or modify it
!! under the terms of the GNU General Public License as published by
!! the Free Software Foundation; either version 3 of the License, or
!! (at your option) any later version.
!!
!! FEMilaro is distributed in the hope that it will be useful, but
!! WITHOUT ANY WARRANTY; without even the implied warranty of
!! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
!! General Public License for more details.
!!
!! You should have received a copy of the GNU General Public License
!! along with FEMilaro; If not, see <http://www.gnu.org/licenses/>.
!!
!! author: Marco Restelli                   <marco.restelli@gmail.com>

!>\brief
!!
!! Collection of Runge-Kutta time integrators
!!
!! \n
!!
!! Provide various RK integrators.
!<----------------------------------------------------------------------
module mod_rk

!-----------------------------------------------------------------------

 use mod_messages, only: &
   mod_messages_initialized, &
   error,   &
   warning, &
   info

 use mod_kinds, only: &
   mod_kinds_initialized, &
   wp

 use mod_state_vars, only: &
   mod_state_vars_initialized, &
   c_stv

 use mod_time_integrators_base, only: &
   mod_time_integrators_base_initialized, &
   c_ode, c_ods, c_tint, t_ti_init_data, t_ti_step_diag

!-----------------------------------------------------------------------
 
 implicit none

!-----------------------------------------------------------------------

! Module interface

 public :: &
   mod_rk_constructor, &
   mod_rk_destructor,  &
   mod_rk_initialized, &
   t_ee, t_hm, t_rk4, t_ssprk54

 private

!-----------------------------------------------------------------------

! Module types

 !> Explicit Euler method
 type, extends(c_tint) :: t_ee
 contains
  procedure, pass(ti) :: init  => ee_init
  procedure, pass(ti) :: step  => ee_step
  procedure, pass(ti) :: clean => ee_clean
 end type t_ee

 !> Heun method
 !!
 !! \f{displaymath}{
 !!  \begin{array}{rcl}
 !!   \displaystyle
 !!   u^{(1)} & = & \displaystyle u^n + \Delta t\,f(u^n) \\[3mm]
 !!   \displaystyle
 !!   u^{n+1} & = & \displaystyle u^n + \frac{\Delta t}{2} \left[
 !!   f(u^n) + f(u^{(1)}) \right].
 !!  \end{array}
 !! \f}
 type, extends(c_tint) :: t_hm
 contains
  procedure, pass(ti) :: init  => hm_init
  procedure, pass(ti) :: step  => hm_step
  procedure, pass(ti) :: clean => hm_clean
 end type t_hm

 !> The <em>classical</em> Runge-Kutta method
 !!
 !! The Butcher tableau for this method is
 !! \f{displaymath}{
 !! \begin{array}{c|cccc}
 !!   0 \\
 !!   \frac{1}{2} & \frac{1}{2} \\[1mm]
 !!   \frac{1}{2} & & \frac{1}{2} \\[1mm]
 !!   1 & & & 1 \\[1mm]
 !!  \hline
 !!   & \frac{1}{6} &\frac{1}{3} &\frac{1}{3} &\frac{1}{6}
 !! \end{array}
 !! \f}
 !! Thanks to the form of this tableau, only one stage must be kept in
 !! memory, using the algorithm:
 !! <ul>
 !!  <li> evaluate \f$k_i\f$ (i.e. the right-hand side)
 !!  <li> set \f$u^{n+1} += b_ik_i\f$
 !!  <li> evaluate \f$k_{i+1}\f$ overwriting \f$k_i\f$.
 !! </ul>
 type, extends(c_tint) :: t_rk4
 contains
  procedure, pass(ti) :: init  => rk4_init
  procedure, pass(ti) :: step  => rk4_step
  procedure, pass(ti) :: clean => rk4_clean
 end type t_rk4

 !> SSPRK54 method
 !!
 !! This is the method discussed in Table A.2 of <a
 !! href="http://dx.doi.org/10.1137/S0036142901389025">[Spiteri,
 !! Ruuth, SIAM J. Num. An., 2002]</a> and <a
 !! href="http://dx.doi.org/10.1007/BF01933264">[Kraaijevanger, BIT
 !! Numerical Mathematics, 1991]</a>.
 !!
 !! In the implementation, we make sure that the consistency condition
 !! \f{displaymath}{
 !! \sum_{k=0}^{i-1} a_{ik} = 1
 !! \f} is verified up to the working precision. Moreover, we notice
 !! that the intermediate time levels can be deduced (as fractions of
 !! \f$\Delta t\f$), setting \f$u^n=0\f$ and \f$f = 1\f$, so that
 !! \f{displaymath}{
 !!  \begin{array}{rclcl}
 !!    t^{(1)} & = & b_{10} & \approx & 0.39175 \\
 !!    t^{(2)} & = & a_{21}\,t^{(1)}+b_{21} & \approx & 0.58608 \\
 !!    t^{(3)} & = & a_{32}\,t^{(2)}+b_{32} & \approx & 0.47454 \\
 !!    t^{(4)} & = & a_{43}\,t^{(3)}+b_{43} & \approx & 0.93501 \\
 !!    t^{(5)} & = & a_{52}\,t^{(2)} + a_{53}\,t^{(3)} +
 !!    a_{54}\,t^{(4)} + b_{53}+b_{54} & = & 1.
 !!  \end{array}
 !! \f}
 !! Notice that \f$b_{54}\f$ is in fact computed using the constraint
 !! \f$t^{(5)} = 1\f$.
 type, extends(c_tint) :: t_ssprk54
 contains
  procedure, pass(ti) :: init  => ssprk54_init
  procedure, pass(ti) :: step  => ssprk54_step
  procedure, pass(ti) :: clean => ssprk54_clean
 end type t_ssprk54




! Module variables

 ! public members
 logical, protected ::               &
   mod_rk_initialized = .false.

 ! private members

 character(len=*), parameter :: &
   this_mod_name = 'mod_rk'

!-----------------------------------------------------------------------

contains

!-----------------------------------------------------------------------

 subroutine mod_rk_constructor()
  character(len=*), parameter :: &
    this_sub_name = 'constructor'

   !Consistency checks ---------------------------
   if( (mod_messages_initialized             .eqv..false.) .or. &
       (mod_kinds_initialized                .eqv..false.) .or. &
       (mod_state_vars_initialized           .eqv..false.) .or. &
       (mod_time_integrators_base_initialized.eqv..false.) ) then
     call error(this_sub_name,this_mod_name, &
                'Not all the required modules are initialized.')
   endif
   if(mod_rk_initialized.eqv..true.) then
     call warning(this_sub_name,this_mod_name, &
                  'Module is already initialized.')
   endif
   !----------------------------------------------

   mod_rk_initialized = .true.
 end subroutine mod_rk_constructor

!-----------------------------------------------------------------------
 
 subroutine mod_rk_destructor()
  character(len=*), parameter :: &
    this_sub_name = 'destructor'
   
   !Consistency checks ---------------------------
   if(mod_rk_initialized.eqv..false.) then
     call error(this_sub_name,this_mod_name, &
                'This module is not initialized.')
   endif
   !----------------------------------------------

   mod_rk_initialized = .false.
 end subroutine mod_rk_destructor

!-----------------------------------------------------------------------
 
 !> Explicit Euler initialization
 subroutine ee_init(ti,ode,dt,t,u0,ods,init_data,step_diag)
  class(c_ode),  intent(in)    :: ode
  real(wp),      intent(in)    :: dt, t
  class(c_stv),  intent(inout) :: u0
  class(c_ods),  intent(inout) :: ods
  class(t_ee),   intent(out)   :: ti
  type(t_ti_init_data), optional, intent(in)  :: init_data
  type(t_ti_step_diag), optional, intent(out) :: step_diag
 
  character(len=*), parameter :: &
    this_sub_name = 'ee_init'

   ti%dt = dt

   ! EE only needs one temporary to store the tendency
   !allocate( ti%work(1) , source=u )
   call u0%source_vect(ti%work,1)
 
 end subroutine ee_init
 
!-----------------------------------------------------------------------
 
 !> Explicit Euler step
 subroutine ee_step(unew,ti,ode,t,unow,ods,step_diag)
  class(t_ee),   intent(inout) :: ti
  class(c_ode),  intent(in)    :: ode
  real(wp),      intent(in)    :: t
  class(c_stv),  intent(in)    :: unow
  class(c_ods),  intent(inout) :: ods
  class(c_stv),  intent(inout) :: unew
  type(t_ti_step_diag), optional, intent(inout) :: step_diag
 
  character(len=*), parameter :: &
    this_sub_name = 'ee_step'

   ! compute the tendency
   call ode%rhs(ti%work(1),t,unow,ods)
   ! compute the time step
   ! unew = unow + dt*tens
   call unew%alt(unow,ti%dt,ti%work(1))
 
 end subroutine ee_step
 
!-----------------------------------------------------------------------
 
 !> Explicit Euler clean-up
 pure subroutine ee_clean(ti,ode,step_diag)
  class(c_ode),  intent(in)    :: ode
  class(t_ee),   intent(inout) :: ti
  type(t_ti_step_diag), optional, intent(inout) :: step_diag

  character(len=*), parameter :: &
    this_sub_name = 'ee_clean'

  deallocate(ti%work)

 end subroutine ee_clean
 
!-----------------------------------------------------------------------

 !> Heun method initialization
 subroutine hm_init(ti,ode,dt,t,u0,ods,init_data,step_diag)
  class(c_ode),  intent(in)    :: ode
  real(wp),      intent(in)    :: dt, t
  class(c_stv),  intent(inout) :: u0
  class(c_ods),  intent(inout) :: ods
  class(t_hm),   intent(out)   :: ti
  type(t_ti_init_data), optional, intent(in)  :: init_data
  type(t_ti_step_diag), optional, intent(out) :: step_diag
 
  character(len=*), parameter :: &
    this_sub_name = 'hm_init'

   ti%dt = dt

   ! HM needs two work arrays
   !allocate( ti%work(2) , source=u )
   call u0%source_vect(ti%work,2)
 
 end subroutine hm_init
 
!-----------------------------------------------------------------------
 
 !> Heun method step
 subroutine hm_step(unew,ti,ode,t,unow,ods,step_diag)
  class(t_hm),   intent(inout) :: ti
  class(c_ode),  intent(in)    :: ode
  real(wp),      intent(in)    :: t
  class(c_stv),  intent(in)    :: unow
  class(c_ods),  intent(inout) :: ods
  class(c_stv),  intent(inout) :: unew
  type(t_ti_step_diag), optional, intent(inout) :: step_diag
 
  character(len=*), parameter :: &
    this_sub_name = 'hm_step'

   ! compute the tendency
   call ode%rhs(ti%work(1),t,unow,ods)

   ! compute the first time step (use unew as temporary)
   call unew%alt(unow,ti%dt,ti%work(1))

   ! compute the second tendency
   call ode%rhs(ti%work(2),t+ti%dt,unew,ods)

   ! compute (twice) the total tendency
   call ti%work(1)%incr(ti%work(2))

   ! compute the complete time step
   call unew%alt(unow,0.5_wp*ti%dt,ti%work(1))

 end subroutine hm_step
 
!-----------------------------------------------------------------------
 
 !> Heun method clean-up
 pure subroutine hm_clean(ti,ode,step_diag)
  class(c_ode),  intent(in)    :: ode
  class(t_hm),   intent(inout) :: ti
  type(t_ti_step_diag), optional, intent(inout) :: step_diag

  character(len=*), parameter :: &
    this_sub_name = 'hm_clean'

  deallocate(ti%work)
 
 end subroutine hm_clean
 
!-----------------------------------------------------------------------

 !> SSPRK54 method initialization
 subroutine ssprk54_init(ti,ode,dt,t,u0,ods,init_data,step_diag)
  class(c_ode),  intent(in)    :: ode
  real(wp),      intent(in)    :: dt, t
  class(c_stv),  intent(inout) :: u0
  class(c_ods),  intent(inout) :: ods
  class(t_ssprk54), intent(out) :: ti
  type(t_ti_init_data), optional, intent(in)  :: init_data
  type(t_ti_step_diag), optional, intent(out) :: step_diag
 
  character(len=*), parameter :: &
    this_sub_name = 'ssprk54_init'

   ti%dt = dt

   ! needs 6 work arrays (this number could be reduced)
   !allocate( ti%work(6) , source=u )
   call u0%source_vect(ti%work,6)
 
 end subroutine ssprk54_init
 
!-----------------------------------------------------------------------
 
 !> SSPRK54 method step
 subroutine ssprk54_step(unew,ti,ode,t,unow,ods,step_diag)
  class(t_ssprk54), intent(inout) :: ti
  class(c_ode),  intent(in)    :: ode
  real(wp),      intent(in)    :: t
  class(c_stv),  intent(in)    :: unow
  class(c_ods),  intent(inout) :: ods
  class(c_stv),  intent(inout) :: unew
  type(t_ti_step_diag), optional, intent(inout) :: step_diag

  real(wp), parameter :: &
    one = 1.0_wp
  real(wp), parameter :: &
    a10 = one,                                             &
    a20 = 0.44437049406734_wp, a21 = one-a20,              &
    a30 = 0.62010185138540_wp, a32 = one-a30,              &
    a40 = 0.17807995410773_wp, a43 = one-a40,              &
    a50 = 0.00683325884039_wp, a52 = 0.51723167208978_wp,  &
          a53 = 0.12759831133288_wp, a54 = one-(a50+a52+a53)
  real(wp), parameter :: &
    b10 = 0.39175222700392_wp, &
    b21 = 0.36841059262959_wp, &
    b32 = 0.25189177424738_wp, &
    b43 = 0.54497475021237_wp, &
    b53 = 0.08460416338212_wp ! b54 defined later, after the ta*
  real(wp), parameter :: &
    ta1 = b10,           &
    ta2 = a21*ta1 + b21, &
    ta3 = a32*ta2 + b32, &
    ta4 = a43*ta3 + b43, &
    ta5 = one
  real(wp), parameter :: &
    b54 = one-(a52*ta2+a53*ta3+a54*ta4+b53)

  character(len=*), parameter :: &
    this_sub_name = 'ssprk54_step'

   ! we use the temporaries as follows:
   !  1:4 stages (stage 5 is stored in unew)
   !  5:6 tendencies

   ! stage 1
   call ode%rhs(ti%work(5),t,unow,ods)
   call ti%work(1)%mlt(a10,unow)
   call ti%work(1)%inlt(ti%dt*b10,ti%work(5))

   ! stage 2
   call ode%rhs(ti%work(5),t+ta1*ti%dt,ti%work(1),ods)
   call ti%work(2)%mlt(a20,unow)
   call ti%work(2)%inlt(a21,ti%work(1))
   call ti%work(2)%inlt(ti%dt*b21,ti%work(5))

   ! stage 3
   call ode%rhs(ti%work(5),t+ta2*ti%dt,ti%work(2),ods)
   call ti%work(3)%mlt(a30,unow)
   call ti%work(3)%inlt(a32,ti%work(2))
   call ti%work(3)%inlt(ti%dt*b32,ti%work(5))

   ! stage 4
   call ode%rhs(ti%work(5),t+ta3*ti%dt,ti%work(3),ods)
   call ti%work(4)%mlt(a40,unow)
   call ti%work(4)%inlt(a43,ti%work(3))
   call ti%work(4)%inlt(ti%dt*b43,ti%work(5))

   ! stage 5
   call ode%rhs(ti%work(6),t+ta4*ti%dt,ti%work(4),ods)
   call unew%mlt(a50,unow)
   call unew%inlv(   (/a52,a53,a54/) , ti%work(2:4) )
   ! Compiler bug: the following line does not work with gfortran:
   ! this should be related to
   ! http://gcc.gnu.org/bugzilla/show_bug.cgi?id=56691
   call unew%inlv( ti%dt*(/b53,b54/) , ti%work(5:6) )

 end subroutine ssprk54_step
 
!-----------------------------------------------------------------------
 
 !> SSPRK54 method celan-up
 pure subroutine ssprk54_clean(ti,ode,step_diag)
  class(c_ode),  intent(in)    :: ode
  class(t_ssprk54), intent(inout) :: ti
  type(t_ti_step_diag), optional, intent(inout) :: step_diag

  character(len=*), parameter :: &
    this_sub_name = 'ssprk54_clean'

  deallocate(ti%work)
 
 end subroutine ssprk54_clean
 
!-----------------------------------------------------------------------

 !> The <em>classical</em> Runge-Kutta method
 subroutine rk4_init(ti,ode,dt,t,u0,ods,init_data,step_diag)
  class(c_ode),  intent(in)    :: ode
  real(wp),      intent(in)    :: dt, t
  class(c_stv),  intent(inout) :: u0
  class(c_ods),  intent(inout) :: ods
  class(t_rk4),  intent(out)   :: ti
  type(t_ti_init_data), optional, intent(in)  :: init_data
  type(t_ti_step_diag), optional, intent(out) :: step_diag

  character(len=*), parameter :: &
    this_sub_name = 'rk4_init'

   ti%dt = dt

   ! needs 6 work arrays (this number could be reduced)
   !allocate( ti%work(2) , source=u )
   call u0%source_vect(ti%work,2)
 
 end subroutine rk4_init
 
!-----------------------------------------------------------------------

 !> The <em>classical</em> Runge-Kutta method
 subroutine rk4_step(unew,ti,ode,t,unow,ods,step_diag)
  class(t_rk4),  intent(inout) :: ti
  class(c_ode),  intent(in)    :: ode
  real(wp),      intent(in)    :: t
  class(c_stv),  intent(in)    :: unow
  class(c_ods),  intent(inout) :: ods
  class(c_stv),  intent(inout) :: unew
  type(t_ti_step_diag), optional, intent(inout) :: step_diag

  real(wp), parameter :: f3=1.0_wp/3.0_wp, f6=1.0_wp/6.0_wp
  real(wp), parameter :: &
    c2=0.5_wp , a21=0.5_wp ,                                   &
    c3=0.5_wp ,              a32=0.5_wp ,                      &
    c4=1.0_wp ,                           a43=1.0_wp ,         &
                 b1=f6     ,  b2=f3     ,  b3=f3     , b4=f6
  character(len=*), parameter :: &
    this_sub_name = 'rk4_step'

   ! stage 1
   call ode%rhs(ti%work(2),t,unow,ods)
   call unew%alt(unow,ti%dt*b1,ti%work(2))

   ! stage 2
   call ti%work(1)%alt(unow,ti%dt*a21,ti%work(2))
   call ode%rhs(ti%work(2),t+c2*ti%dt,ti%work(1),ods)
   call unew%inlt(ti%dt*b2,ti%work(2))

   ! stage 3
   call ti%work(1)%alt(unow,ti%dt*a32,ti%work(2))
   call ode%rhs(ti%work(2),t+c3*ti%dt,ti%work(1),ods)
   call unew%inlt(ti%dt*b3,ti%work(2))

   ! stage 4
   call ti%work(1)%alt(unow,ti%dt*a43,ti%work(2))
   call ode%rhs(ti%work(2),t+c4*ti%dt,ti%work(1),ods)
   call unew%inlt(ti%dt*b4,ti%work(2))

 end subroutine rk4_step

!-----------------------------------------------------------------------
 
 pure subroutine rk4_clean(ti,ode,step_diag)
  class(c_ode),  intent(in)    :: ode
  class(t_rk4),  intent(inout) :: ti
  type(t_ti_step_diag), optional, intent(inout) :: step_diag

  character(len=*), parameter :: &
    this_sub_name = 'rk4_clean'

  deallocate(ti%work)
 
 end subroutine rk4_clean
 
!-----------------------------------------------------------------------

end module mod_rk

